﻿#ifndef MODEL_H
#define MODEL_H

#include <stack>
#include <assimp/Importer.hpp>
#include <assimp/scene.h>
#include <assimp/postprocess.h>

#include "Mesh.h"
#include "Shader.h"
#include "stb_image/stb_image.h"

class Model
{
public:
	explicit Model(const char *path, const Material &material, const bool gamma = false)
		: material(material), gammaCorrection(gamma)
	{
		loadModel(path);
	}

	void Draw() const
	{
		for (const auto &mesh: meshes)
			mesh.Draw(material);
	}

private:
	std::vector<Texture> textures_loaded; // 存储到已加载的所有纹理，优化以确保纹理不会被加载多次。
	std::vector<Mesh> meshes;
	const Material &material;
	std::string directory;

	bool gammaCorrection;

	void loadModel(const std::string &path)
	{
		Assimp::Importer import;
		const aiScene *scene = import.ReadFile(path, aiProcess_Triangulate | aiProcess_FlipUVs);

		if (!scene || scene->mFlags & AI_SCENE_FLAGS_INCOMPLETE || !scene->mRootNode)
		{
			std::cout << "ERROR::ASSIMP::" << import.GetErrorString() << std::endl;
			return;
		}
		directory = path.substr(0, path.find_last_of('/'));

		processNode(scene->mRootNode, scene);
	}

	void processNode(const aiNode *node, const aiScene *scene)
	{
		// 使用显式栈替代递归，将递归改为迭代方式处理节点树，避免潜在的栈溢出问题。
		std::stack<const aiNode *> nodeStack;
		nodeStack.push(node);

		while (!nodeStack.empty())
		{
			const aiNode *currentNode = nodeStack.top();
			nodeStack.pop();

			// 读取节点所有的网格
			for (unsigned int i = 0; i < currentNode->mNumMeshes; i++)
			{
				const aiMesh *mesh = scene->mMeshes[currentNode->mMeshes[i]];
				meshes.push_back(processMesh(mesh, scene));
			}

			// 将子节点压入栈中，反转以保持原有顺序
			for (int i = static_cast<int>(currentNode->mNumChildren) - 1; i >= 0; i--)
			{
				nodeStack.push(currentNode->mChildren[i]);
			}
		}
	}

	static Mesh processMesh(const aiMesh *mesh, const aiScene *scene)
	{
		std::vector<Vertex> vertices;
		std::vector<unsigned int> indices;
		// std::vector<Texture> textures;

		for (unsigned int i = 0; i < mesh->mNumVertices; i++)
		{
			Vertex vertex;
			auto vec2Zero = glm::vec2(0.0f, 0.0f);
			auto vec3Zero = glm::vec3(0.0f, 0.0f, 0.0f);

			// 处理顶点位置、法线和纹理坐标
			glm::vec3 tempVector3;
			tempVector3.x = mesh->mVertices[i].x;
			tempVector3.y = mesh->mVertices[i].y;
			tempVector3.z = mesh->mVertices[i].z;
			vertex.Position = tempVector3;

			if (mesh->HasNormals())
			{
				tempVector3.x = mesh->mNormals[i].x;
				tempVector3.y = mesh->mNormals[i].y;
				tempVector3.z = mesh->mNormals[i].z;
				vertex.Normal = tempVector3;
			}
			else
			{
				vertex.Normal = vec3Zero;
			}

			// 一个顶点最多可以包含8个不同的纹理坐标。
			// 这里假设不是多纹理坐标模型，因此使用第[0]组。
			if (mesh->mTextureCoords[0])
			{
				// 纹理坐标
				glm::vec2 vec;
				vec.x = mesh->mTextureCoords[0][i].x;
				vec.y = mesh->mTextureCoords[0][i].y;
				vertex.TexCoords = vec;

				if (mesh->HasTangentsAndBitangents())
				{
					// 切线
					tempVector3.x = mesh->mTangents[i].x;
					tempVector3.y = mesh->mTangents[i].y;
					tempVector3.z = mesh->mTangents[i].z;
					vertex.Tangent = tempVector3;

					// 副切线
					tempVector3.x = mesh->mBitangents[i].x;
					tempVector3.y = mesh->mBitangents[i].y;
					tempVector3.z = mesh->mBitangents[i].z;
					vertex.Bitangent = tempVector3;
				}
				else
				{
					vertex.Tangent = vec3Zero;
					vertex.Bitangent = vec3Zero;
				}
			}
			else
			{
				vertex.TexCoords = vec2Zero;
			}

			vertices.push_back(vertex);
		}

		// 处理索引
		// 遍历每个网格的面，并获取对应的顶点索引。
		for (unsigned int i = 0; i < mesh->mNumFaces; i++)
		{
			const aiFace &face = mesh->mFaces[i];
			for (unsigned int j = 0; j < face.mNumIndices; j++)
				indices.push_back(face.mIndices[j]);
		}

		// 处理材质
		// aiMaterial *ai_material = scene->mMaterials[mesh->mMaterialIndex];

		// N:序号
		// diffuse: texture_diffuseN
		// specular: texture_specularN
		// normal: texture_normalN

		// 1. diffuse maps
		// std::vector<Texture> diffuseMaps = loadMaterialTextures(ai_material, aiTextureType_DIFFUSE, "texture_diffuse");
		// textures.insert(textures.end(), diffuseMaps.begin(), diffuseMaps.end());
		// // 2. specular maps
		// std::vector<Texture> specularMaps = loadMaterialTextures(ai_material, aiTextureType_SPECULAR, "texture_specular");
		// textures.insert(textures.end(), specularMaps.begin(), specularMaps.end());
		// // 3. normal maps
		// std::vector<Texture> normalMaps = loadMaterialTextures(ai_material, aiTextureType_HEIGHT, "texture_normal");
		// textures.insert(textures.end(), normalMaps.begin(), normalMaps.end());
		// // 4. height maps
		// std::vector<Texture> heightMaps = loadMaterialTextures(ai_material, aiTextureType_AMBIENT, "texture_height");
		// textures.insert(textures.end(), heightMaps.begin(), heightMaps.end());
		//
		// material.textures = textures;

		return {vertices, indices};
	}

	std::vector<Texture> loadMaterialTextures(const aiMaterial *mat, const aiTextureType type,
	                                          const std::string &typeName)
	{
		std::vector<Texture> textures;
		for (unsigned int i = 0; i < mat->GetTextureCount(type); i++)
		{
			aiString str;
			mat->GetTexture(type, i, &str);

			// 检查纹理是否已经加载，如果是，则直接跳到下一次迭代
			bool skip = false;
			for (auto &j: textures_loaded)
			{
				if (std::strcmp(j.path.data(), str.C_Str()) == 0)
				{
					textures.push_back(j);
					skip = true;
					break;
				}
			}
			if (skip) continue;

			Texture texture;
			const bool isGammaCorrected = typeName == "texture_diffuse" ? gammaCorrection : false; // 对diffuse启用gamma校正
			texture.id = TextureFromFile(str.C_Str(), this->directory, isGammaCorrected);
			texture.type = typeName;
			texture.path = str.C_Str();
			textures.push_back(texture);
			textures_loaded.push_back(texture);
		}
		return textures;
	}

	static unsigned int TextureFromFile(const char *path, const std::string &directory, const bool gamma)
	{
		auto filename = std::string(path);
		filename = directory + '/' + filename;

		unsigned int textureID;
		glGenTextures(1, &textureID);

		int width, height, nrComponents;
		if (unsigned char *data = stbi_load(filename.c_str(), &width, &height, &nrComponents, 0))
		{
			GLint format;
			if (nrComponents == 1)
				format = GL_RED;
			else if (nrComponents == 3)
				format = gamma ? GL_SRGB : GL_RGB;
			else
				format = gamma ? GL_SRGB_ALPHA : GL_RGBA;

			glBindTexture(GL_TEXTURE_2D, textureID);
			glTexImage2D(GL_TEXTURE_2D, 0, format, width, height, 0, format, GL_UNSIGNED_BYTE, data);
			glGenerateMipmap(GL_TEXTURE_2D);

			glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);
			glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
			glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR);
			glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);

			stbi_image_free(data);
		}
		else
		{
			std::cout << "内部贴图加载失败， 检测路径: " << path << std::endl;
			stbi_image_free(data);
		}

		return textureID;
	}
};
#endif //MODEL_H
